import hdbscan
import seaborn as sns
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from pyproj import Transformer


def convert_to_utm(df):
    """WGS84转UTM Zone 50N（与Java代码相同的EPSG:32650）"""
    transformer = Transformer.from_crs("EPSG:4326", "EPSG:32650", always_xy=True)
    lon = df['longitude'].astype(float).values
    lat = df['latitude'].astype(float).values
    easting, northing = transformer.transform(lon, lat)
    return easting, northing


# 直接在内存中处理轨迹数据
def process_and_cluster_trajectory(input_path):
    # 读取原始数据（保持与Java相同的字段索引）
    df = pd.read_csv(input_path, header=0, encoding='gbk')
    df['easting'], df['northing'] = convert_to_utm(df)

    # 初始化模型参数
    clusterer = hdbscan.HDBSCAN(
        min_cluster_size=10,  # 最小簇大小
        min_samples=5,  # 核心点最小邻居数
        cluster_selection_method='eom'  # 簇选择方法
    )

    # 转换为HDBSCAN可处理的numpy数组
    # 选择空间特征：东坐标、北坐标、高度（三维空间聚类）
    test_data = df[['easting', 'northing', 'height']].values

    # 执行聚类
    cluster_labels = clusterer.fit_predict(test_data)

    # 将聚类结果添加回DataFrame
    df['Cluster'] = cluster_labels

    # 输出聚类轨迹条数和噪声点数量
    unique_clusters, counts = np.unique(cluster_labels, return_counts=True)
    for cluster, count in zip(unique_clusters, counts):
        if cluster == -1:
            print(f"噪声点数量: {count}")
        else:
            print(f"簇 {cluster} 中的轨迹条数: {count}")

    # 可视化聚类结果
    plt.figure(figsize=(10, 6))
    palette = sns.color_palette('husl', len(set(cluster_labels)) - (1 if -1 in cluster_labels else 0))
    cluster_colors = [palette[col] if col >= 0 else (0.7, 0.7, 0.7) for col in cluster_labels]
    plt.scatter(test_data[:, 0], test_data[:, 1], c=cluster_colors, s=30, alpha=0.9)
    plt.title("HDBSCAN聚类结果")
    plt.show()

    return df


# 使用函数处理并聚类轨迹数据
df_processed = process_and_cluster_trajectory('../src/main/resources/track250224_foshan.csv')